#define DEBUG_TYPE "bssa"

#include "flowTracking.h"

using namespace llvm;

std::vector<BasicBlock *> ProcessedBB;

static cl::opt<string> InputFilename("xmlfile", cl::Required, cl::desc(
		"Informe um arquivo XML valido."), cl::value_desc("fileXML"));

//Receive a predicate and include it on predicate attribute
Pred::Pred(Value *p) {
	predicate = p;
}

//Receive a instruction and include it on insts vector
void Pred::addInst(Instruction *i) {
	insts.push_back(i);
}

//Receive a function pointer and include it on funcs vector
void Pred::addFunc(Function *f) {
	funcs.push_back(f);
}

//Return the total instruction count
int Pred::getNumInstrucoes() {
	return (insts.size());
}

int Pred::getNumFunctions() {
	return (funcs.size());
}

//Return the predicate
Value *Pred::getPred() {
	return (predicate);
}

//Return the instruction stored on insts vector, pointed for parameter i
Instruction *Pred::getInst(int i) {
	if (i < (signed int) insts.size())
		return (insts[i]);
	else
		return NULL;
}

Function *Pred::getFunc(int i) {
	if (i < (signed int) funcs.size())
		return (funcs[i]);
	else
		return NULL;
}

//Return true of *op instruction is gated (if it is stored on insts vector) for the predicate
bool Pred::isGated(Instruction *op) {
	unsigned int i;

	for (i = 0; i < insts.size(); i++) {
		if (op == insts[i])
			return true;
	}

	return false;
}

////////////////////////////////////////////////////////////////////////////////////////////////////
//Passes which are used by bSSA pass
void bSSA::getAnalysisUsage(AnalysisUsage &AU) const {

	AU.addRequired<PostDominatorTree> ();
	AU.addRequired<moduleDepGraph> ();

	// This pass will modifies the program, but not the CFG
	AU.setPreservesCFG();

}

//For each module, this method is executed
bool bSSA::runOnModule(Module &M) {

	// Add input file


	parsersXML::parserXML *xmlparser = new parsersXML::parserXML();
	if (!xmlparser->setFile(InputFilename)) {
		errs() << "Não foi possível carregar o arquivo XML." << "\n";
		;
		return false;
	}
	std::vector<parsersXML::parserXML::funcaoFonte> funcoes_fontes;
	std::vector<std::string> funcoes_saida;
	funcoes_fontes = xmlparser->getFontes();
	funcoes_saida = xmlparser->getSaidas();

	moduleDepGraph& DepGraph = getAnalysis<moduleDepGraph> ();
	Function *F;
	//Getting dependency graph
	Graph *g = DepGraph.depGraph;

	//src stores the instructions which are source of secret information
	//dst stores the instructions  which are public output channel like printf()
	std::vector<Value *> src, dst;
	std::vector<BasicBlock *> srcB;
	std::vector<bool> srcUM;
	std::vector<StringRef> srcFileName, dstFileName;
	std::vector<unsigned> srcLine, dstLine;

	for (Module::iterator Mit = M.begin(), Mend = M.end(); Mit != Mend; ++Mit) {
		F = Mit;
		// Iterate over all Basic Blocks of the Function
		for (Function::iterator Fit = F->begin(), Fend = F->end(); Fit != Fend; ++Fit) {
			makeTable(Fit, F); //Creating in memory the table with predicates and gated instructions
		}

	}

	//Including control edges into dependence graph.
	incGraph(g);

	Instruction *A;
	StringRef File;
	MDNode *N;

	//Interates on all source code in order to get the sources of address (secret information) and sinks (instructions like printf)
	for (Module::iterator F = M.begin(), eM = M.end(); F != eM; ++F) {
		for (Function::iterator BB = F->begin(), e = F->end(); BB != e; ++BB) {


			//errs() << "Basic block (name=" << BB->getName() << ") has "
            // << BB->size() << " instructions.\n";

             //errs() << "Operações: \n";
			BasicBlock *fBB = &F->front();
			for (BasicBlock::iterator I = BB->begin(), ee = BB->end(); I != ee; ++I) {

				errs() << *I << " ---- \n";

				//Getting Line and Filename of each source and destination
				A = cast<Instruction> (I);
				if (MDNode * N = A->getMetadata("dbg")) {
					DILocation Loc(N);
					File = Loc.getFilename();
				} else {
					File = "Unknown Source File";
				}
				N = A->getMetadata("dbg");
				DILocation Loc(N);
				unsigned Line = Loc.getLineNumber();

				if (dyn_cast<Instruction> (I)->getOpcode()
						== Instruction::PtrToInt) {
					src.push_back(I);
					srcLine.push_back(Line);
					srcFileName.push_back(File);
				} 
				//else if (dyn_cast<Instruction> (I)->getOpcode()
				// 		== Instruction::Alloca) {
				// 	src.push_back(I);
				// 	srcLine.push_back(Line);
				// 	srcFileName.push_back(File);
				// }

				//If is a function call
				if (CallInst *CI = dyn_cast<CallInst>(I)) {
					Function *Callee = CI->getCalledFunction();
					if (Callee) {
						StringRef Name = Callee->getName();
						//errs() << Name << "\n";
						// verifica se a função é uma função de destino
						for (std::vector<string>::iterator it =
								funcoes_saida.begin(); it
								!= funcoes_saida.end(); ++it) {
							if (Name.equals(*it)) {
								dst.push_back(I);
								dstLine.push_back(Line);
								dstFileName.push_back(File);
								break;
							}
						}

						// verifica se a função é uma função de entrada
						// for (std::vector<parsersXML::parserXML::funcaoFonte>::iterator it =
						// 		funcoes_fontes.begin(); it
						// 		!= funcoes_fontes.end(); ++it) {
						for(std::vector<int>::size_type i = 0; i != funcoes_fontes.size(); i++) {
							std::string cs = funcoes_fontes[i].getNome();
							const char * c = cs.c_str();
							if (Name.equals(c)) {
								srcB.push_back(fBB);
								src.push_back(I);
								srcUM.push_back(funcoes_fontes[i].usaMemoria());
								srcLine.push_back(Line);
								srcFileName.push_back(File);
								break;
							}
						}

					}
				}
			}
		}
	}

	//Stats about sources and sinks
	errs() << "********* Flow Tracking Summary ************";
	errs() << "\n\nSources: " << src.size() << "\n";
	errs() << "Public channels: " << dst.size() << "\n";

	unsigned int countWarning = 0;

	//Search leaks for each (source, sink). If there is a leak, include the tainted subgraph into set of tainted subgraphs.
	int c = 0;
	for (unsigned int i = 0; i < src.size(); i++) {
		for (unsigned int j = 0; j < dst.size(); j++) {

			GraphNode* source = g->findOpNode(src[i]);

			std::set<GraphNode*> visitedNodes;
			//errs() << "\n\nNodos anteriores a: " << source->getName() << " " << source->getLabel() << " " <<srcUM[i] << "\n";


			g->dfsVisitBackMemo(source, visitedNodes);
			
			std::vector<Value *> srcV;
			std::vector<StringRef> srcFileNameV;
			std::vector<unsigned> srcLineV;
			int verifica = false;

			Instruction *A;
			StringRef File;
			MDNode *N;
			

			if(srcUM[i]){
				for (std::set<GraphNode*>::iterator it = visitedNodes.begin(); it != visitedNodes.end(); ++it) {
	               			 //errs() << (*it)->getName() << " " << (*it)->getLabel() << " " <<   "\n";
	                		if((*it)->getLabel() == "Call malloc" || (*it)->getLabel() == "alloca"){
						for (BasicBlock::iterator I = srcB[i]->begin(), ee = srcB[i]->end(); I != ee; ++I) {
							
							StringRef Name;
							if (CallInst *CI = dyn_cast<CallInst>(I)) {
								Function *Callee = CI->getCalledFunction();
								if (Callee) {
									Name = Callee->getName();
								}
							}
							/******************/			
							A = cast<Instruction> (I);
							if (MDNode * N = A->getMetadata("dbg")) {
								DILocation Loc(N);
								File = Loc.getFilename();
							} else {
								File = "Unknown Source File";
							}
							N = A->getMetadata("dbg");
							DILocation Loc(N);
							unsigned Line = Loc.getLineNumber();
							/******************/	
							errs() << "\n\nNodos enderecos: " << (*I) << " -x- " << (*it)->getLabel() <<  "\n";						
							if (dyn_cast<Instruction> (I)->getOpcode() == Instruction::Alloca || Name.equals("malloc")) {
								//errs() << "\n\nNodos enderecos: " << (*I) << " " << (*it)->getLabel() <<  "\n";

								srcV.push_back(I);
								srcLineV.push_back(Line);
								srcFileNameV.push_back(File);								
							}
						}
						for (unsigned int k = 0; k < srcV.size(); k++) {
							Graph subGV = g->generateSubGraph(srcV[k], src[i]);

							Graph::iterator gItV = subGV.begin();
							Graph::iterator gIteV = subGV.end();

							if(gItV != gIteV){
								errs() << "\n\nEntrou ";
								src[i] = srcV[k];
								srcLine[i] = srcLineV[k];
								srcFileName[i] = srcFileNameV[k];
								verifica = true;
								break;	
																
							}
						}
						
	                		}
	        		}
			}
			

			Graph subG = g->generateSubGraph(src[i], dst[j]);

			Graph::iterator gIt = subG.begin();
			Graph::iterator gIte = subG.end();
			if (gIt != gIte || (gIt != gIte && verifica)) {
				countWarning++;
				ostringstream ss, ss0, ss1;
				ss0 << "./subgraphLines" << c << ".dot";
				ss << "./subgraphLLVM" << c << ".dot";
				c++;
				//ss1 << srcFileName[i].data();
				ss1 << "file " << srcFileName[i].data() << " line "
						<< srcLine[i] << " to file " << dstFileName[j].data()
						<< " line " << dstLine[j] << " ";
				subG.toDot(ss1.str(), ss.str()); //Make one file .dot for each tainted subgraph.
				subG.toDotLines(ss1.str(), ss0.str());
			}

		}
	}

	//Número de warnings
	errs() << "Number of address leak paths " << countWarning << "\n\n";

	//Make the file .dot including the dependence graph
	std::string Filename = "/tmp/fullGraph.dot";
	//Print dependency graph (in dot format);
	g->toDot("Grafo", Filename);

	return false;
}

//Increase graph including control edges
void bSSA::incGraph(Graph *g) {
	unsigned int i;
	int j;

	//For all predicates in predicatesVector
	for (i = 0; i < predicatesVector.size(); i++) {
		//Locates the predicate (icmp instrution) Localiza o predicado (instrução icmp) from the graph
		GraphNode *predNode = g->findNode(predicatesVector[i]->getPred());

		//For each predicate, iterates on the list of gated INSTRUCTIONS
		for (j = 0; j < predicatesVector[i]->getNumInstrucoes(); j++) {
			GraphNode *instNode = g->findNode(predicatesVector[i]->getInst(j));
			if (predNode != NULL && instNode != NULL) {//If the instruction is on the graph, make a edge
				g->addEdge(predNode, instNode, etControl);
			}
		}

		//For each predicate, iterates on the list of gated FUNCTIONS
		for (j = 0; j < predicatesVector[i]->getNumFunctions(); j++) {
			Function *F = predicatesVector[i]->getFunc(j);
			//For each function, iterates on its basic blocks
			for (Function::iterator Fit = F->begin(), Fend = F->end(); Fit
					!= Fend; ++Fit) {
				//For each basic block, iterates on its instructions
				for (BasicBlock::iterator bBIt = Fit->begin(), bBEnd =
						Fit->end(); bBIt != bBEnd; ++bBIt) {
					GraphNode *instNode = g->findNode(bBIt);
					if (predNode != NULL && instNode != NULL)
						g->addEdge(predNode, instNode, etControl);
				}
			}
		}

	}

}

//It receives a BasicBLock and makes table of predicates and its respective gated instructions
void bSSA::makeTable(BasicBlock *BB, Function *F) {
	Value *condition;
	TerminatorInst *ti = BB->getTerminator();
	BranchInst *bi = NULL;
	SwitchInst *si = NULL;

	PostDominatorTree &PD = getAnalysis<PostDominatorTree> (*F);

	ProcessedBB.clear();
	if ((bi = dyn_cast<BranchInst> (ti)) && bi->isConditional()) { //If the terminator instruction is a conditional branch
		condition = bi->getCondition();
		//Including the predicate on the predicatesVector
		predicatesVector.push_back(new Pred(condition));
		//Make a "Flooding" on each sucessor gated the instruction on Influence Region of the predicate
		for (unsigned int i = 0; i < bi->getNumSuccessors(); i++) {
			findIR(BB, bi->getSuccessor(i), PD);
		}
	} else if ((si = dyn_cast<SwitchInst> (ti))) {
		condition = si->getCondition();
		//Including the predicate on the predicatesVector
		predicatesVector.push_back(new Pred(condition));
		//Make a "Flooding" on each sucessor gated the instruction on Influence Region of the predicate
		for (unsigned int i = 0; i < si->getNumSuccessors(); i++) {
			findIR(BB, si->getSuccessor(i), PD);
		}
	}
}

//Flooding until reach a posdominator node
void bSSA::findIR(BasicBlock *bBOring, BasicBlock *bBSuss,
		PostDominatorTree &PD) {

	Pred *p;
	TerminatorInst *ti = bBSuss->getTerminator();

	//If the basic block has been processed, do not advance
	for (unsigned int x = 0; x < ProcessedBB.size(); x++) {
		if (ProcessedBB[x] == bBSuss) {
			return;
		}
	}

	//Including the basic block in the processed vector
	ProcessedBB.push_back(bBSuss);

	//If the basic block is a posdominator and is not the start basic block, just gate the PHI instructions
	if (PD.dominates(bBSuss, bBOring) && bBSuss != bBOring) {
		//Find PHI instructions
		for (BasicBlock::iterator bBIt = bBSuss->begin(), bBEnd = bBSuss->end(); bBIt
				!= bBEnd; ++bBIt) {
			if (dyn_cast<Instruction> (bBIt)->getOpcode() == Instruction::PHI) {
				//if there is a PHI's argument gated, gate the PHI instruction
				p = predicatesVector.back();
				for (unsigned int k = 0; k
						< dyn_cast<Instruction> (bBIt)->getNumOperands(); k++) {
					if (p->isGated(dyn_cast<Instruction> (
							dyn_cast<Instruction> (bBIt)->getOperand(k)))) {
						p->addInst(bBIt);
						break;
					}
				}
			}
		}
		return;
	} else { //Advance the flooding
		//Instruction will be gated whit the bBOring predicate
		for (BasicBlock::iterator bBIt = bBSuss->begin(), bBEnd = bBSuss->end(); bBIt
				!= bBEnd; ++bBIt) {
			p = predicatesVector.back();

			//If is a function call which is defined on the same module
			if (CallInst *CI = dyn_cast<CallInst>(&(*bBIt))) {
				Function *F = CI->getCalledFunction();
				if (F != NULL)
					if (!F->isDeclaration() && !F->isIntrinsic()) {
						gateFunction(F, p);
					}
			}

			//Gate the other instructions
			p->addInst(bBIt);
		}
		//If there is successor, go there
		for (unsigned int i = 0; i < ti->getNumSuccessors(); i++) {
			findIR(bBOring, ti->getSuccessor(i), PD);
		}
	}

}

//All instrutions of function F are gated with predicate p
void bSSA::gateFunction(Function *F, Pred *p) {
	// Iterate over all Basic Blocks of the Function
	//for (Function::iterator Fit = F->begin(), Fend = F->end(); Fit != Fend; ++Fit) {
	//for (BasicBlock::iterator bBIt = Fit->begin(), bBEnd = Fit->end(); bBIt != bBEnd; ++bBIt) {
	p->addFunc(F);
	//}
	//}
}

//Print all predicates and its respective gated instructions
void bSSA::printGate() {
	for (unsigned int i = 0; i < predicatesVector.size(); i++) {
		errs() << "\n\n" << predicatesVector[i]->getPred()->getName() << "\n";
		for (int j = 0; j < predicatesVector[i]->getNumInstrucoes(); j++) {
			errs() << predicatesVector[i]->getInst(j)->getOpcodeName() << "\n";

		}
	}

}

char bSSA::ID = 0;
static RegisterPass<bSSA> X("flowTracking",
		"Implicit and Explicit flow analysis");

